import numpy as np
import matplotlib.pyplot as plt
import math
import random
#Code for figure
#create function

def generate_matrix(n, m, r, elements):
    U, _, _ = np.linalg.svd(np.random.rand(n,r))
    _, _, V = np.linalg.svd(np.random.rand(r,m))
    A = np.dot(U[:,:r], np.dot(np.diag(elements), V[:m,:r].T))
    return A


def watch(A,deg):
    B =  np.hstack((np.ones(int(len(A)*len(A[1])*deg)),np.zeros(int(len(A)*len(A[1])*(1-deg)))))
    random.shuffle(B)
    B = np.reshape(B,[len(A),len(A[1])])
    return B


def f(X):
    return np.linalg.norm(np.multiply(X,Obs)-A_obs,ord = 'fro')**2+0.05*np.linalg.norm(X,ord = 'nuc')
#原函数，输入n维list，输出1维list

def grad(X):
    U1,S1,V1 = np.linalg.svd(X,full_matrices=0)
    D= np.dot(2,np.multiply(X,Obs)-A_obs)+np.dot(0.05,np.dot(U1,V1))
    D = np.array(D)
    return D
#梯度，输入n维list,输出n维list

def compare(x1,x2):
    return x1 if f(x1)<=f(x2) else x2
#比较函数，输入两个n维list,输出一个n维list
def gradw(x,rate,y):
    return [grad(x)[i]+((x[i]-y[i])/math.sqrt(rate)) for i in range(len(grad(x)))]
def Pr(x,y,iter,rate):
    for i in range(iter - 1):
        x = x - np.dot(rate , gradw(x,rate,y))
    return x
def GNSA(origin,epochs,rate,p):
    x = np.array(origin)
    y = np.array(origin)
    z = np.array(origin)
    dot_set = [x]
    for i in range(epochs - 1):
        y = (1 - co(p, i)) * x + co(p, i) * z
        x = compare(y - rate * grad(y),x- rate*grad(x))
        z = z - (rate / (co(p, i))) * grad(y)
        dot_set.append(x)
    return dot_set
#GNSA，输出epochs*n 维列表
def GD(origin,epochs,rate):
    x= np.array(origin)
    dot_set = [x]
    for i in range(epochs-1):
        x=x-rate*grad(x)
        dot_set.append(x)
    return dot_set
#梯度下降，输入初值，步数，学习率，输出epochs*n 维列表
def PROX(origin,epochs,rate,p):
    x = np.array(origin)
    y = np.array(origin)
    z = np.array(origin)
    dot_set = [x]
    for i in range(epochs - 1):
        y = (1 - co(p, i)) * x + co(p, i) * z
        x = Pr(x,y,iter=10,rate=rate)
        z = z - (rate / (co(p, i))) * grad(x)
        dot_set.append(x)
    return dot_set
def co(p,t):
    return p/(t+p)
def NSA(origin,epochs,rate,p):
    x= np.array(origin)
    y= np.array(origin)
    z= np.array(origin)
    dot_set = [x]
    for i in range(epochs-1):
        y=(1-co(p,i))*x+co(p,i)*z
        x=y-rate*grad(y)
        z=z-(rate/(co(p,i)))*grad(y)
        dot_set.append(x)
    return dot_set
#Nesterov梯度下降，输出epochs*n 维列表

def pr_track(y):
    return
def pr_fun_value(y,min):
    plt.semilogy([i for i in range(len(y))],[f(y[i])-min for i in range(len(y))])
    return
#绘制轨迹图与函数值下降图
#调试
np.random.seed(4)
n1= 250
m1= 100
rank = 7
de=0.2
A=generate_matrix(n=n1,m=m1,r=rank,elements=[1,2,3,4,5,6,7])
Obs=watch(A,de)
A_obs = np.multiply(A,Obs)
or1 = np.zeros([len(A),len(A[1])])
or2 = np.ones([len(A),len(A[1])])
print(grad(or1))
end=600
cend=800
e1=0.1
#print(grad_descent(origin=[-2,3,5],epochs=end,rate=0.01)[end-1])#梯度下降调试
#print(NSA(origin=[-2,3,5],epochs=end,rate=0.01,p=3)[end-1])#Nesterov梯度下降测试
#print(GNSA(origin=[-2,3,5],epochs=end,rate=0.01,p=3)[end-1])
#min1 = f(GNSA(origin=or1,epochs=cend,rate=e1,p=3)[cend-1])
cp = GNSA(origin=or1,epochs=cend,rate=e1,p=3)
min1 = np.min([f(cp[i]) for i in range(cend) ])
pr_fun_value(GD(origin=or1,epochs=end,rate=e1),min1)
pr_fun_value(NSA(origin=or1,epochs=end,rate=e1,p=3),min1)
pr_fun_value(PROX(origin=or1,epochs=end,rate=e1,p=3),min1)
plt.legend(loc='best',frameon=True, labels=['GD','NAG','PROX'])
plt.xlabel('iteration')
plt.ylabel('f - f*')
plt.show()